import cv2
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib as mpl
from tensorflow.keras.layers import Input
from src.yolo3.model import *
from src.yolo3.detect import *
from src.utils.image import *
from src.utils.datagen import *
from src.utils.fixes import *
fix_tf_gpu()
# Importing all necessary libraries
import cv2
import os
# Read the video from specified path
vid = cv2.VideoCapture(r"C:\Users\Niharika\Downloads\emotions.mp4")
try:
# creating a folder named data
if not os.path.exists('data'):
os.makedirs('data')
# if not created then raise error
except OSError:
print('Error: Creating directory of data')
# frame
currentframe = 0
while (True):
# reading from frame
success, frame = vid.read()
if success:
# continue creating images until video remains
name = './data/frame' + str(currentframe) + '.jpg'
print('Creating...' + name)
# writing the extracted images
cv2.imwrite(name, frame)
# increasing counter so that it will
# show how many frames are created
currentframe += 1
else:
break
# Release all space and windows once done
vid.release()
cv2.destroyAllWindows()
def prepare_model(approach):
'''
Prepare the YOLO model
'''
global input_shape, class_names, anchor_boxes, num_classes, num_anchors, model
# shape (height, width) of the imput image
input_shape = (416, 416)
# class names
if approach == 1:
class_names = ['W']
elif approach == 2:
class_names = ['H', 'V', 'W']
elif approach == 3:
class_names = ['W','WH','WV','WHV']
else:
raise NotImplementedError('Approach should be 1, 2, or 3')
# anchor boxes
if approach == 1:
anchor_boxes = np.array(
[
np.array([[ 76, 59], [ 84, 136], [188, 225]]) /32, # output-1 anchor boxes
np.array([[ 25, 15], [ 46, 29], [ 27, 56]]) /16, # output-2 anchor boxes
np.array([[ 5, 3], [ 10, 8], [ 12, 26]]) /8 # output-3 anchor boxes
],
dtype='float64'
)
else:
anchor_boxes = np.array(
[
np.array([[ 73, 158], [128, 209], [224, 246]]) /32, # output-1 anchor boxes
np.array([[ 32, 50], [ 40, 104], [ 76, 73]]) /16, # output-2 anchor boxes
np.array([[ 6, 11], [ 11, 23], [ 19, 36]]) /8 # output-3 anchor boxes
],
dtype='float64'
)
# number of classes and number of anchors
num_classes = len(class_names)
num_anchors = anchor_boxes.shape[0] * anchor_boxes.shape[1]
# input and output
input_tensor = Input( shape=(input_shape[0], input_shape[1], 3) ) # input
num_out_filters = ( num_anchors//3 ) * ( 5 + num_classes ) # output
# build the model
model = yolo_body(input_tensor, num_out_filters)
# load weights
weight_path = r'C:\Users\Niharika\Desktop\IVA\weights\pictor-ppe-v302-a2-yolo-v3-weights.h5'
model.load_weights( weight_path )
def get_detection(img):
# save a copy of the img
act_img = img.copy()
# shape of the image
ih, iw = act_img.shape[:2]
# preprocess the image
img = letterbox_image(img, input_shape)
img = np.expand_dims(img, 0)
image_data = np.array(img) / 255.
# raw prediction from yolo model
prediction = model.predict(image_data)
# process the raw prediction to get the bounding boxes
boxes = detection(
prediction,
anchor_boxes,
num_classes,
image_shape = (ih, iw),
input_shape = (416,416),
max_boxes = 10,
score_threshold=0.3,
iou_threshold=0.45,
classes_can_overlap=False)
# convert tensor to numpy
boxes = boxes[0].numpy()
# draw the detection on the actual image
return draw_detection(act_img, boxes, class_names)
def plt_imshow(img):
plt.figure(figsize=(5, 5))
plt.imshow(img)
plt.axis('off')
Model detects only workers.
# prepare the model
prepare_model(approach=1)
for i in range(11):
# read the image
img = cv2.imread( f'C:/Users/Niharika/Desktop/IVA/extras/sample-images/{i}.jpg' )
# resize
img = letterbox_image(img, input_shape)
# get the detection on the image
img = get_detection(img)
# show the image
plt_imshow(img[:, :, ::-1])
Model detects worker, hat, and vest (three object classes) individually.
# prepare the model
prepare_model(approach=2)
for i in range(11):
# read the image
img = cv2.imread( f'C:/Users/Niharika/Desktop/IVA/extras/sample-images/{i}.jpg' )
# resize
img = letterbox_image(img, input_shape)
# get the detection on the image
img = get_detection(img)
# show the image
plt_imshow(img[:, :, ::-1])
Model localizes workers in the input image and directly classifies each detected worker as W, WH, WV, or WHV.
# prepare the model
prepare_model(approach=3)
for i in range(11):
# read the image
img = cv2.imread( f'C:/Users/Niharika/Desktop/IVA/extras/sample-images/{i}.jpg' )
# resize
img = letterbox_image(img, input_shape)
# get the detection on the image
img = get_detection(img)
# show the image
plt_imshow(img[:, :, ::-1])
# prepare the model
prepare_model(approach=3)
# read the image
img = cv2.imread( r'C:\Users\Niharika\Desktop\IVA\extras\sample-images\11.jpg' )
# resize
img = letterbox_image(img, input_shape)
# get the detection on the image
img = get_detection(img)
# show the image
plt_imshow(img[:, :, ::-1])